import asyncio
import json
import os
from pathlib import Path

from dotenv import load_dotenv
from langchain.agents import initialize_agent, AgentType
from langchain_deepseek import ChatDeepSeek
from langchain_mcp_adapters.client import MultiServerMCPClient

CONFIG_FILE = Path("weather_mcp_config.json")


async def main():
    cfg = json.loads(CONFIG_FILE.read_text(encoding="utf-8"))

    # 1. 实例化客户端（不再用 async with）
    client = MultiServerMCPClient(cfg["mcpServers"])

    # 2. 取工具列表
    tools = await client.get_tools()

    # 3. 挂到 Agent
    # 读取env配置信息
    load_dotenv(override=True)
    deepseek_api_key = os.getenv("DEEPSEEK_API_KEY")

    # 初始化 deepseek
    pop_llm = ChatDeepSeek(
        model="deepseek-chat",
        temperature=0,
        max_tokens=None,
        timeout=None,
        max_retries=2,
        api_key=deepseek_api_key,
    )
    agent = initialize_agent(
        tools=tools,
        llm=pop_llm,
        agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
    )

    # 4. 使用
    ans = await agent.arun("北京现在天气如何？")
    print(ans)

    # 5. 手动关闭（可选，进程结束时会自动清理）
    # await client.close()


if __name__ == "__main__":
    asyncio.run(main())
